#!/usr/bin/env python2
# coding=utf-8
import rospy
# import pandas as pd
import matplotlib.pyplot as plt
from px4_swarm_control.msg import Swarm_Position
from px4_swarm_control.msg import remote_uwb_info
from px4_swarm_control.msg import sp_target_position
from geometry_msgs.msg import Point
from std_msgs.msg import Header
import time
import numpy as np
import os
import random
from math import sin
from math import cos
from math import sqrt
from time import sleep


class Swarm:
    def __init__(self, num=4):
        rospy.init_node('position_controler', anonymous=True)
        self.id = self.loadParam("id", 0)
        self.control_rate = self.loadParam("rate", 30)
        self.rate = rospy.Rate(self.control_rate)
        self.path = os.path.abspath('..')
        self.agent_num = num
        self.size = self.agent_num
        self.agent_safe_r = self.loadParam("safe_r", 0.5)
        self.global_vel = np.zeros((self.agent_num, 2))
        self.target_position = [0, 0]
        self.safe_move_area_x = [self.loadParam("safe_x_min", 0),
                                 self.loadParam("safe_x_max", 5)]
        self.safe_move_area_y = [self.loadParam("safe_y_min", 0),
                                 self.loadParam("safe_y_max", 5)]
        self.v_max = self.loadParam("max_vel", 1)
        self.step = 1 / self.control_rate  # time
        self.min_v = 0.001
        self.global_position = np.zeros((self.agent_num, 2))
        self.PI = 3.1415926
        # diff with algorithm set self id adj_m = 0
        self.adj_matrix = np.ones((self.size, self.size))
        for i in range(len(self.adj_matrix)):
            self.adj_matrix[self.id][i] = 0
            self.adj_matrix[i][self.id] = 0
        self.swarm_position_sub = rospy.Subscriber("/swarm_position_id" + str(self.id), Swarm_Position,
                                                   self.swarm_position_callback)
        self.swarm_status_sub = rospy.Subscriber("/swarm_state/remote_nodes", remote_uwb_info,
                                                 self.swarm_status_callback)
        self.swarm_target_pub = rospy.Publisher("/global_drone_target_position_id" + str(self.id), sp_target_position,
                                                queue_size=1)
        self.swarm_status = remote_uwb_info()
        self.last_swarm_status = remote_uwb_info()
        for i in range(self.size):
            self.adj_matrix[i][i] = 0
        self.get_position = 1

    def set_adj_matrix(self, set_adj_matirx):
        self.adj_matrix = set_adj_matirx.copy()

    def swarm_status_callback(self, msg):
        self.swarm_status = msg
        # 两次false 确认后才进行更新
        if len(self.swarm_status.active) == len(self.last_swarm_status.active):
            for i in range(len(self.swarm_status.active)):
                self.swarm_status.active[i] = self.swarm_status.active[i] or self.last_swarm_status.active[i]
        self.id = msg.self_id
        self.update_adj_matrix()
        self.last_swarm_status = msg
        # rospy.loginfo_throttle(0.5, "ID: [" + self.swarm_status.self_id +
        #                        "] current active : " + str(self.swarm_status.active))

    def loadParam(self, name, default):
        name = rospy.get_name() + "/" + name
        res = rospy.get_param(name, default)
        rospy.loginfo("loded param " + name + " :" + str(res))
        return res

    # 如果当前节点不活跃 把该节点的邻接矩阵置0,活跃的话置1
    def update_adj_matrix(self):
        for i in range(len(self.adj_matrix)):
            self.adj_matrix[self.id][i] = 0
            self.adj_matrix[i][self.id] = 0
        # 一个id的时候 基本用于测试，开启邻接矩阵
        if len(self.swarm_status.active) == 0:
            for i in range(len(self.adj_matrix)):
                self.adj_matrix[self.id][i] = 1
                self.adj_matrix[i][self.id] = 1

        # 一旦收到active信息，严格按照其进行更新邻接矩阵
        if len(self.swarm_status.node_ids) <= self.agent_num and len(self.swarm_status.active) > 0:
            node_id = 0
            node_id_active = 0
            for i in range(len(self.swarm_status.node_ids)):
                node_id
                node_id_active = self.swarm_status.active[i]
                if not node_id_active:
                    self.adj_matrix[self.id][node_id] = 0
                    self.adj_matrix[node_id][self.id] = 0
                if node_id_active and self.adj_matrix[self.id][node_id] == 0:
                    self.adj_matrix[self.id][node_id] = 1
                    self.adj_matrix[node_id][self.id] = 1
                    rospy.loginfo(
                        "active status changed node active: " + str(node_id))

        self.update_toplogic()

    # 更新拓扑图连接
    def update_toplogic(self):
        pass

    def control(self):
        self.update_adj_matrix()
        while not rospy.is_shutdown():
            self.rate.sleep()
            if not self.get_position:
                continue
            vx, vy = self.control_by_index(self.id)
            self.sp_publish()
        rospy.spin()

    def sp_publish(self):
        sp_msg = sp_target_position()
            sp_msg.header.stamp = rospy.Time.now()
            if self.get_position:
                sp_msg.global_sp_pose.x = self.global_position[self.id][0] + \
                    vx * 10 * self.step
                sp_msg.global_sp_pose.y = self.global_position[self.id][1] + \
                    vy * 10 * self.step
                sp_msg.global_sp_pose_v.x = vx
                sp_msg.global_sp_pose_v.y = vy
            else:
                sp_msg.global_sp_pose.x = self.global_position[self.id][0]
                sp_msg.global_sp_pose.y = self.global_position[self.id][1]
                sp_msg.global_sp_pose_v.x = 0
                sp_msg.global_sp_pose_v.y = 0
            rospy.loginfo_throttle(
                0.5, "[" + str(self.id) + "]calculated vx: " + str(vx) + " vy: " + str(vy))
            self.swarm_target_pub.publish(sp_msg)

    def swarm_position_callback(self, msg):
        temp = 0
        self.global_position = np.zeros((self.agent_num, 2))
        self.global_vel = np.zeros((self.agent_num, 2))
        for i in range(self.agent_num):
            self.global_position[i][0] = msg.swarm_pose[i].x
            self.global_position[i][1] = msg.swarm_pose[i].y
            self.global_vel[i][0] = msg.swarm_vel[i].x
            self.global_vel[i][1] = msg.swarm_vel[i].y
            temp = temp + \
                self.global_position[i][0] + self.global_position[i][0]
        if abs(temp) > 0.001:
            self.get_position = 1

    def get_target_position(self):
        v = self.control_by_index(self.id)
        self.target_position[0] = self.global_position[self.id][0] + v[0]
        self.target_position[1] = self.global_position[self.id][1] + v[1]

    def get_distance(self, target_position, index):
        res = np.zeros(self.size)
        for i in range(self.size):
            res[i] = res[i] + sqrt((target_position[i][0] - target_position[index][0]) *
                                   (target_position[i][0] - target_position[index][0]) +
                                   (target_position[i][1] - target_position[index][1]) *
                                   (target_position[i][1] - target_position[index][1]))
        return res

    def control_by_index(self, index):
        x = self.global_position[index][0]
        y = self.global_position[index][1]
        if x < self.safe_move_area_x[0]:
            print "x out of safe"
            return [self.v_max, 0]
        if x > self.safe_move_area_x[1]:
            print "x out of safe"
            return [-self.v_max, 0]
        if y < self.safe_move_area_y[0]:
            print "y out of safe"
            return [0, self.v_max]
        if y > self.safe_move_area_y[1]:
            print "y out of safe"
            return [0, -self.v_max]
        vx = 0
        vy = 0
        cur_distance = self.get_distance(self.global_position, index)
        not_safe = 1
        while not_safe:
            not_safe = 0
            target_position = self.global_position.copy()

            for i in range(self.size):
                dx = target_position[i][0] - target_position[index][0]
                dy = target_position[i][1] - target_position[index][1]
                if self.adj_matrix[index][i] != 1:
                    continue
                if 2 * self.agent_safe_r > cur_distance[i] > 0:
                    vx = vx + 5 * -dx * pow((2 * self.agent_safe_r - sqrt(dx * dx + dy * dy)), 2) / sqrt(
                        dx * dx + dy * dy)
                    vy = vy + 5 * -dy * pow((2 * self.agent_safe_r - sqrt(dx * dx + dy * dy)), 2) / sqrt(
                        dx * dx + dy * dy)
                    continue
                vx = vx + dx
                vy = vy + dy
            if sqrt(vx * vx + vy * vy) > self.v_max:
                _vx = self.v_max * vx / sqrt(vx * vx + vy * vy)
                _vy = self.v_max * vy / sqrt(vx * vx + vy * vy)
                vx = _vx
                vy = _vy

            target_position[index][0] = target_position[index][0] + \
                vx * self.step
            target_position[index][1] = target_position[index][1] + \
                vy * self.step
            dist = self.get_distance(target_position, index)
            for i in range(self.size):
                if dist[i] > self.agent_safe_r * 2:
                    not_safe = 0
            if not_safe:
                vx = vx / 4
                vy = vy / 4
                if abs(vx) < self.min_v:
                    vx = 0
                if abs(vy) < self.min_v:
                    vy = 0
                if abs(vx) < self.min_v and abs(vy) < self.min_v:
                    return [0, 0]
            else:
                break
        return [vx, vy]
        pass


if __name__ == "__main__":
    drone_swarm = Swarm()
    drone_swarm.control()
